fix: report which half of a JAX first call is stalled - #1529

Closes part of #1528 (phase 3 of the jax-compile-stall epic). Library leg — merges first; the two workspace PRs consume it.

The problem

The first call to a jax.jit/vmap/grad wrapper contains two waits — trace/lower/compile, then jax.block_until_ready execution — and the instrumentation could not tell a stalled run which one it was in.

Two defects, both of which made a stall report the wrong diagnosis.

1. The heartbeat hardcoded "still compiling". A stalled CI run therefore logged JAX jit still compiling ... 1770s elapsed while its stack sat in jax.block_until_ready. That is positive evidence for the wrong cause, and five quarantine markers across the two test workspaces were written against it, every one calling this an "intermittent XLA compile stall". It is not one — compilation completes, in ~12–18s, and execution never returns.

2. The compile/execute split was emitted only after BOTH halves finished. A stalled run never finishes the second, so it reported neither number and the split characterised only the healthy case. Twenty stalled runs could not say that compilation had in fact completed.

The change

The compile half is logged the instant func() returns. The heartbeat then names the materialize half it has moved into and repeats the compile time on every beat — a stalled run is SIGKILLed and reaches no summary line, so whatever beat last got to stderr has to carry the split by itself.

Also extracts _block_until_ready() so the execution half is substitutable in tests without JAX installed (the library suite is numpy-only, so the real path cannot produce a hang), and so the pre-existing broad except Exception guards only the optional import jax rather than swallowing failures from the wait itself.

No call site changes — all four (Fitness._vmap/_jit/_grad, analysis/latent) inherit this through log_on_first_compile.

This is what found the bug

From autolens_workspace_test run 33076408637, verbatim:

13:33:33  JAX jit compiled vectorized (vmap) likelihood function in 15.7s
          and is still waiting for the result to materialize, 270s elapsed...

On the old code that identical run reads still compiling ... 270s elapsed and produces a sixth marker with the wrong name. The campaign's root cause was found in the hours after this landed.

Tests

Five new, all asserting relationships rather than wordings — the trap #1517 recorded, where its own tests could not catch a 300s dump default colliding with a 300s cap:

• the compile line lands before the execution half is waited on
• it is present while a run is still parked in execution
• every execute-phase beat names that phase and carries the compile time
• no beat still claims to be compiling once past the boundary
• the real _block_until_ready stays a no-op when JAX is absent

test_jax_compile.py 23 passed. Full suite 2224 passed, 3 skipped.

Note for anyone running the suite locally: -n auto is unusable in this repo — parametrised test IDs embed object reprs with memory addresses, so xdist workers disagree on collection and error out. Pre-existing, unrelated to this change.

Heart

Not consulted — pyauto-heart is unreachable from the web-github environment, as recorded on phases 1 and 2 of this epic.

Generated by Claude Code